import numpy as np
import pyworld
import glob
from hparams import hparams
import librosa
import os

def feature_world(wav,para):
    fs = para.fs
    wav = wav.astype(np.float64)
    f0, timeaxis = pyworld.harvest(wav, fs, frame_period=para.frame_period, f0_floor=71.0, f0_ceil=800.0)
    
    sp = pyworld.cheaptrick(wav, f0, timeaxis, fs)
    ap = pyworld.d4c(wav, f0, timeaxis, fs)
    coded_sp =pyworld.code_spectral_envelope(sp, fs, para.coded_dim)
    return f0,timeaxis,sp,ap,coded_sp
    
 
def processing_wavs(file_wavs,para):
    
    f0s = []
    coded_sps = []
    for file in file_wavs:
        print("processing %s"%(file))
        # 读取音频文件
        fs = para.fs
        wav, _ = librosa.load(file, sr=fs, mono=True)
        
        # 提取world 特征,采集f0和coded_sp
        f0,_,_,_,coded_sp=feature_world(wav,para)
        print(coded_sp.shape)
        f0s.append(f0)
        coded_sps.append(coded_sp)
        
    
    # 计算log_f0的 均值和std
    log_f0s = np.ma.log(np.concatenate(f0s))
    log_f0s_mean = log_f0s.mean()
    log_f0s_std = log_f0s.std()
    
    # 计算 coded_sp 的均值和 标准差
    coded_sps_array = np.concatenate(coded_sps,axis=0)  # coded_sp的维度  T * D
    coded_sps_mean = np.mean(coded_sps_array,axis=0,keepdims = True)
    coded_sps_std = np.std(coded_sps_array,axis=0,keepdims = True)

    # 利用 coded_sp 的均值和 标准差 对特征进行正则
    coded_sps_norm = []
    for coded_sp in coded_sps:
        coded_sps_norm.append(  (coded_sp- coded_sps_mean)/ coded_sps_std )
        
    return log_f0s_mean,log_f0s_std,coded_sps_mean,coded_sps_std,coded_sps_norm
    
    
if __name__ == "__main__":
    
    para = hparams()
    print(para.path_train_wavs)
    # 遍历所有 spks    
    for spk in para.spk_list:
        print("processing features for %s"%(spk_id))
        # 获取每个spk的wav文件存放路径 
        dir_train = os.path.join(para.path_train_wavs,spk)
        wavs = glob.glob(dir_train+'/*wav') 
        f0_mean,f0_std,mecp_mean,mecp_std, mecps = processing_wavs(wavs,para)
        
        # 获取保存路径
        path_save = os.path.join(para.path_catch_feas,spk_id)
        os.makedirs(path_save,exist_ok = True)
        
        # 进行数据保存
        np.save(os.path.join(path_save,'static_f0.npy'),np.array([f0_mean,f0_std],dtype=object))
        np.save(os.path.join(path_save,'static_mecp.npy'),np.array([mecp_mean,mecp_std],dtype=object))
        np.save(os.path.join(path_save,'data.npy'),np.array(mecps,dtype=object))
        
        
        
        
        
    
    
    
    
    
    
    
    